use std::fmt::{self, Debug};
use std::num::{NonZeroUsize};
use std::sync::atomic::{self};
use diff;
use frag;

// Traits

/// Trait for "abstract" nodes of the tree that do not correspond directly to underlying fragtree.
pub trait Abstract: Sized {
  type State;
  type Context: NodeContext;

  fn initial_state(&self, ctx: &mut Self::Context) -> Self::State;
  fn render(&self, ctx: &mut Self::Context, state: &Self::State) -> Virt<Self>;

  fn compatible(&self, other: &Self) -> bool;
  fn should_update(&self,
                   state: &Self::State,
                   next_cfg: &Self,
                   next_state: &Self::State) -> bool;
}

/// Trait for context info passed along to all nodes.
pub trait NodeContext {
  fn push_path(&mut self, el: PathElem);
  fn pop_path(&mut self);

  fn push_new_uid(&mut self, uid_gen: &mut UidGen) -> Uid {
    let uid = uid_gen.next();
    self.push_path(PathElem::ByUid(uid));
    uid
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Data types

///
/// Virtual node -- specifies the desired shape of a node. Virtual nodes can be "imposed" onto live
/// nodes, which updates the live node to conform to the specification given by the virtual node.
#[derive(Clone, PartialEq, Eq)]
pub enum Virt<P> {
  Frag(frag::Frag),
  Abstract(P),
  Branch(VirtBranch<P>),
}

#[derive(Clone, PartialEq, Eq)]
pub struct VirtBranch<P> {
  spine: frag::FragSpine,
  info: frag::FragInfo,
  arms: Vec<Virt<P>>,
}

///
/// Live node -- specifies the current shape of a node with some additional state attached to nodes.
#[derive(Debug, PartialEq, Eq)]
pub enum Live<P, S> {
  Frag,
  Abstract(Box<AbstractState<P, S>>),
  Branch(LiveBranch<P, S>),
}

#[derive(Debug)]
pub struct LiveBranch<P, S> {
  spine: frag::FragSpine,
  arms: Vec<Live<P, S>>,
  arm_uids: Vec<Uid>,
  uid_gen: UidGen,
}

#[derive(Debug, PartialEq, Eq)]
pub struct AbstractState<P, S> {
  props: P,
  state: S,
  rendered: Live<P, S>,
}

/// Unique identifiers for tagging arms of `Live::Branch`.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub struct Uid(NonZeroUsize);

/// `Uid` generator.
pub struct UidGen(atomic::AtomicUsize);

/// Generic structure for result of comparing trees. The reason it has so many type parameters is
/// that the values within differ greatly based on if the comparison takes ownership or not. See the
/// `LiveVirtCmp` and `LiveVirtCmpMut` variations below.
#[derive(Debug, Clone, Eq, PartialEq)]
enum Cmp<D, F, B1, B2, A1, A2> {
  Different(D),          // different spines
  Frag(F),               // both trees are Frag; RHS has fragment F
  Branch(B1, B2),        // both trees are Branch; RHS has info I
  Abstract(A1, A2),      // both trees are Abstract
}

type LiveVirtCmp<'a, P, S> =
  Cmp<(),                            // different
      &'a frag::Frag,                // frag
      &'a LiveBranch<P, S>,          // branch
      &'a VirtBranch<P>,             // ... (branch)
      &'a AbstractState<P, S>,       // abstract
      &'a P>;                        // ... (abstract)

type LiveVirtCmpMut<'a, P, S> =
  Cmp<Virt<P>,                       // different
      frag::Frag,                    // frag
      &'a mut LiveBranch<P, S>,      // branch
      VirtBranch<P>,                 // ... (branch)
      &'a mut AbstractState<P, S>,   // abstract
      P>;                            // ... (abstract)

/// Paths used for locating `Abstract` nodes given by instructions on which arms to crawl in order
/// to reach the desired node.

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum PathElem {
  Direct,
  ByUid(Uid),
}

pub type PathBuf = Vec<PathElem>;
pub type Path = [PathElem];

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum FollowFailed {
  UidNotPresent,
  WrongElemType,
  WrongNodeType,
}

impl fmt::Display for FollowFailed {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "path does not match any node in the tree")
  }
}

impl std::error::Error for FollowFailed {}


////////////////////////////////////////////////////////////////////////////////////////////////////

// ! (never type) and () (unit type) can be used for trees without abstract nodes.

impl Abstract for ! {
  type State = ();
  type Context = ();
  fn initial_state(&self, _: &mut ()) {}
  fn render(&self, _: &mut (), _state: &()) -> Virt<!> { *self }
  fn compatible(&self, _other: &Self) -> bool { *self }
  fn should_update(&self, _: &(), _: &Self, _: &()) -> bool { false }
}

impl NodeContext for () {
  fn push_path(&mut self, _: PathElem) {}
  fn pop_path(&mut self) {}
}

// Uid and UidGen

impl Debug for UidGen {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "UidGen(..)") }
}

impl Uid {
  pub fn dummy() -> Self {
    Uid(unsafe { NonZeroUsize::new_unchecked(1) })
  }
}

impl UidGen {
  pub fn new() -> Self {
    UidGen(atomic::AtomicUsize::new(2))
  }

  fn next(&mut self) -> Uid {
    let n = self.0.fetch_add(1, atomic::Ordering::Relaxed);
    Uid(unsafe { NonZeroUsize::new_unchecked(n) })
  }
}

// TODO: remove this -- it (was) only used for hardcoding paths for debugging
impl From<usize> for PathElem {
  fn from(x: usize) -> PathElem {
    PathElem::ByUid(Uid(NonZeroUsize::new(x).expect("Uid(0) is invalid")))
  }
}

// ignore uids when comparing Live nodes

impl<P, S> PartialEq for LiveBranch<P, S>
where P: PartialEq, S: PartialEq {
  fn eq(&self, rhs: &Self) -> bool {
    self.spine == rhs.spine
      && self.arms == rhs.arms
  }
}

impl<P, S> Eq for LiveBranch<P, S> where P: Eq, S: Eq {}


////////////////////////////////////////////////////////////////////////////////////////////////////

// Constructors

impl<P> Virt<P> {
  pub fn frag(frg: frag::Frag) -> Self {
    Virt::Frag(frg)
  }

  pub fn branch(spine: frag::FragSpine,
                info: frag::FragInfo,
                arms: Vec<Virt<P>>) -> Self {
    Virt::Branch(VirtBranch { spine, arms, info })
  }

  pub fn abs_props(props: P) -> Self {
    Virt::Abstract(props)
  }

  pub fn info_mut(&mut self) -> Option<&mut frag::FragInfo> {
    match *self {
      Virt::Branch(VirtBranch { ref mut info, .. }) => Some(info),
      _ => None,
    }
  }
}

impl<P, S> Live<P, S> {
  pub fn frag() -> Self {
    Live::Frag
  }
}

// Querying Live nodes: follow, compare

impl<P, S> Live<P, S> {

  /// Follow the given path to find a live node, and the underlying node it corresponds to.
  pub fn follow_underlying_mut<'s, 'p, 'u>
    (&'s mut self, path: &'p Path, underly: &'u mut frag::Frag)
     -> Result<(&'s mut Live<P, S>, &'u mut frag::Frag), FollowFailed> {
    let mut cur_underly = underly;
    let mut cur_live = self;
    for &path_el in path.iter() {
      match path_el {
        PathElem::Direct =>
          match *cur_live {
            Live::Abstract(ref mut abs) => cur_live = &mut abs.rendered,
            _ => return Err(FollowFailed::WrongElemType),
          }

        PathElem::ByUid(uid) =>
          match *cur_live {
            Live::Branch(ref mut branch) => {
              let idx = branch.find_arm_index(uid)?;
              cur_live = &mut branch.arms[idx];
              cur_underly = &mut cur_underly[idx];
            }
            _ => return Err(FollowFailed::WrongElemType),
          }
      }
    }
    Ok((cur_live, cur_underly))
  }

  /// Follow the given path to find a live node.
  pub fn follow(&self, path: &Path) -> Result<&Live<P, S>, FollowFailed> {
    let mut cur_live = self;
    for &path_el in path.iter() {
      match path_el {
        PathElem::Direct =>
          match *cur_live {
            Live::Abstract(ref abs) => cur_live = &abs.rendered,
            _ => return Err(FollowFailed::WrongElemType),
          }

        PathElem::ByUid(uid) =>
          match *cur_live {
            Live::Branch(ref branch) => {
              let idx = branch.find_arm_index(uid)?;
              cur_live = &branch.arms[idx];
            }
            _ => return Err(FollowFailed::WrongElemType),
          }
      }
    }
    Ok(cur_live)
  }

  /// Follow the given path to find an abstract live node.
  pub fn follow_abstract(&self, path: &Path) -> Result<&AbstractState<P, S>, FollowFailed> {
    match self.follow(path)? {
      Live::Abstract(ref abs) => Ok(abs),
      _ => Err(FollowFailed::WrongNodeType),
    }
  }

  /// Compare trees shallowly.
  fn compare<'a>(&'a self, virt: &'a Virt<P>) -> LiveVirtCmp<'a, P, S> {
    match *self {
      Live::Branch(ref b) =>
        match *virt {
          Virt::Branch(ref vb) => Cmp::Branch(b, vb),
          _ => Cmp::Different(()),
        }

      Live::Frag =>
        match *virt {
          Virt::Frag(ref u) => Cmp::Frag(u),
          _ => Cmp::Different(()),
        }

      Live::Abstract(ref abs) =>
        match *virt {
          Virt::Abstract(ref abs_node) => Cmp::Abstract(abs, abs_node),
          _ => Cmp::Different(()),
        }
    }
  }

  /// Compare trees shallowly (mutable; takes ownership of RHS).
  fn compare_mut(&mut self, virt: Virt<P>) -> LiveVirtCmpMut<P, S> {
    match *self {
      Live::Branch(ref mut b) =>
        match virt {
          Virt::Branch(vb) => Cmp::Branch(b, vb),
          _ => Cmp::Different(virt),
        }

      Live::Frag =>
        match virt {
          Virt::Frag(u) => Cmp::Frag(u),
          _ => Cmp::Different(virt),
        }

      Live::Abstract(ref mut abs) =>
        match virt {
          Virt::Abstract(abs_node) => Cmp::Abstract(abs, abs_node),
          _ => Cmp::Different(virt),
        }
    }
  }
}

impl<P, S> LiveBranch<P, S> {
  pub fn find_arm_index(&self, uid: Uid) -> Result<usize, FollowFailed> {
    self.arm_uids
        .iter()
        .position(|&arm_uid| arm_uid == uid)
        .ok_or(FollowFailed::UidNotPresent)
  }
}

// Imposing virtual tree onto live tree + fragtree.

impl<P, S> Live<P, S> where P: Abstract<State=S> {

  /// Impose a virtual tree onto this live tree, updating this tree's structure, and the given
  /// underlying tree's structure, so that it matches the virtual tree.
  pub fn impose(&mut self,
                virt: Virt<P>,
                underly: &mut frag::Frag,
                ctx: &mut P::Context) {
    match self.compare_mut(virt) {
      Cmp::Branch(mut b, vb) =>
        diff_branches(&mut b, vb, underly, ctx),

      Cmp::Abstract(abs, props) =>
        if abs.props.compatible(&props) {
          abs.set_props(props, underly, ctx);
        } else {
          self.impose_overwrite(Virt::Abstract(props), underly, ctx);
        }

      Cmp::Frag(frg) =>
        self.impose_overwrite(Virt::Frag(frg), underly, ctx),

      Cmp::Different(virt) =>
        self.impose_overwrite(virt, underly, ctx),
    }
  }

  fn impose_overwrite(&mut self,
                      virt: Virt<P>,
                      underly: &mut frag::Frag,
                      ctx: &mut P::Context) {
    // TODO: unmount() logic here, in the future
    let (new_live, new_underly) = virt.build(ctx);
    *self = new_live;
    *underly = new_underly;
  }
}

fn diff_branches<'a, 'c, P, S>(branch: &'a mut LiveBranch<P, S>,
                               v_branch: VirtBranch<P>,
                               base_underlying: &'a mut frag::Frag,
                               context: &'c mut P::Context)
where P: Abstract<State=S> {
  // update branch spines and underlying branch info
  if branch.spine != v_branch.spine {
    branch.spine = v_branch.spine.clone();
    base_underlying.set_spine(v_branch.spine);
  }
  base_underlying.set_info(v_branch.info);

  // perform diff algorithm on the branch arms (live vs. virt), yielding a "patch"
  let patch = diff::gen(&mut branch.arms, v_branch.arms);

  // apply the patch to both the live branch and underlying node
  diff::ApplyPatch::apply_patch(
    &mut TreeApplyPatch { branch, base_underlying, context },
    patch,
  ).unwrap()
  // there's no reason this should unwrap() fail -- the patch is applied to the same data that
  // generated it.
}

impl<P> diff::Cmp<Virt<P>> for Live<P, P::State> where P: Abstract {
  fn similar(&self, virt: &Virt<P>) -> bool {
    match self.compare(virt) {
      Cmp::Frag(_) => true,
      Cmp::Branch(b1, b2) => b1.spine == b2.spine,
      Cmp::Abstract(abs, n) => abs.props.compatible(n),
      Cmp::Different(()) => false,
    }
  }
}

/// Helper struct for implementing diff::ApplyPatch. Patches applied to this structure will be
/// applied to both the live branch and the underlying tree. In particular, `overwrite()` will call
/// `.impose()` rather than naively overwriting nodes in the tree.
struct TreeApplyPatch<'a, 'c, P, S, CTX> {
  branch: &'a mut LiveBranch<P, S>,
  base_underlying: &'a mut frag::Frag,
  context: &'c mut CTX,
}

impl<'a, 'c, P> diff::ApplyPatch<Virt<P>> for TreeApplyPatch<'a, 'c, P, P::State, P::Context>
where P: Abstract {
  fn len(&self) -> usize {
    self.branch.arms.len()
  }

  fn remove(&mut self, i: usize) {
    // TODO: keep track of arms/arm_uids invariants somewhere else, e.g. a LiveBranch<> impl.
    self.branch.arms.remove(i);
    self.branch.arm_uids.remove(i);
    self.base_underlying.remove(i);
  }

  fn insert(&mut self, i: usize, v: Virt<P>) {
    let uid = self.context.push_new_uid(&mut self.branch.uid_gen);
    let (live, under) = v.build(self.context);
    self.context.pop_path();
    self.branch.arms.insert(i, live);
    self.branch.arm_uids.insert(i, uid);
    self.base_underlying.insert(i, under);
  }

  fn overwrite(&mut self, i: usize, v: Virt<P>) -> bool {
    let uid = self.branch.arm_uids[i];
    let live = &mut self.branch.arms[i];
    self.context.push_path(PathElem::ByUid(uid));
    live.impose(v, &mut self.base_underlying[i], self.context);
    self.context.pop_path();
    true
  }
}

impl<P> Virt<P> where P: Abstract {
  /// Build an actual tree (live and underlying) from this virtual tree. The resulting
  /// trees are in sync.
  fn build(self, ctx: &mut P::Context) -> (Live<P, P::State>, frag::Frag) {
    match self {
      Virt::Frag(frg) =>
        (Live::Frag, frg),

      Virt::Branch(VirtBranch { spine, arms, info }) => {
        let mut underlying_arms = Vec::with_capacity(arms.len());
        let mut branch = LiveBranch {
          spine: spine.clone(),
          arms: Vec::with_capacity(arms.len()),
          arm_uids: Vec::with_capacity(arms.len()),
          uid_gen: UidGen::new(),
        };
        for v_arm in arms {
          let uid = ctx.push_new_uid(&mut branch.uid_gen);
          let (live, under) = v_arm.build(ctx);
          ctx.pop_path();
          branch.arms.push(live);
          branch.arm_uids.push(uid);
          underlying_arms.push(under);
        }
        (Live::Branch(branch),
         frag::Frag::new(spine, info, underlying_arms))
      }

      Virt::Abstract(props) => {
        let state = props.initial_state(ctx);
        let virt = props.render(ctx, &state);
        ctx.push_path(PathElem::Direct);
        let (rendered, under) = virt.build(ctx);
        ctx.pop_path();
        (Live::Abstract(Box::new(AbstractState { props, state, rendered })),
         under)
       }
    }
  }
}

// Manipulating abstract node state

impl<P, S> AbstractState<P, S> {
  pub fn props(&self) -> &P { &self.props }
  pub fn state(&self) -> &S { &self.state }
}

impl<P> AbstractState<P, P::State>
where P: Abstract {

  pub fn set_props(&mut self, props: P, underly: &mut frag::Frag, ctx: &mut P::Context) {
    if self.props.should_update(&self.state, &props, &self.state) {
      self.props = props;
      self.rerender(underly, ctx)
    }
  }

  /*
  pub fn set_state(&mut self, next_state: A::State, underly: &mut U) {
    if self.node.should_update(&self.state, &self.node, &next_state) {
      self.state = next_state;
      self.rerender(underly);
    }
  }
   */

  pub fn rerender(&mut self, underly: &mut frag::Frag, ctx: &mut P::Context) {
    let virt_rendered = self.props.render(ctx, &self.state);
    ctx.push_path(PathElem::Direct);
    self.rendered.impose(virt_rendered, underly, ctx);
    ctx.pop_path();
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// debugging

impl<P, S> Live<P, S> where S: Debug {
  pub fn print_structure(&self) {
    self.print_structure_aux(0, "".into())
  }

  fn print_structure_aux(&self, indent: usize, prefix: String) {
    for _ in 0..indent { print!(" "); }
    print!("{}", prefix);

    match *self {
      Live::Frag =>
        print!("frag\n"),

      Live::Branch(ref b) => {
        print!("branch {:?}\n", b.spine);
        for (l, uid) in b.arms.iter().zip(b.arm_uids.iter()) {
          l.print_structure_aux(2 + indent, format!("[{:?}] ", uid));
        }
      }

      Live::Abstract(ref a) => {
        print!("abstract (state={:?})\n", a.state);
        a.rendered.print_structure_aux(2 + indent, "".into());
      }
    }
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod test {
  use super::*;
  use diff;
  use std::ops::{Index, IndexMut};

  #[derive(Clone, Eq, PartialEq, Debug)]
  struct StrT {
    spine: String,
    children: Vec<StrT>,
  }

  type L = Live<StrT, !>;
  type V = Virt<StrT, !>;

  /// Helper to succinctly create trees.
  fn parse<B, A, I, T>(i: &mut I, build: &B, alt: &A) -> Option<T>
  where I: Iterator<Item=char>,
        B: Fn(String, Vec<T>) -> T,
        A: Fn() -> T {
    let mut spine = String::from("");
    loop {
      match i.next().unwrap() {
        ')' => return None,
        '*' => return Some(alt()),
        '(' => break,
        c => spine.push(c),
      }
    }
    let mut arms = Vec::new();
    loop {
      match parse(i, build, alt) {
        Some(e) => arms.push(e),
        None => return Some(build(spine, arms)),
      }
    }
  }

  fn tree(s: &str) -> StrT {
    parse(&mut s.chars(), &|spine, children| StrT { spine, children }, &|| panic!(".")).unwrap()
  }

  fn live(s: &str) -> L {
    parse(&mut s.chars(), &|spine, arms| Branch { spine, arms }.into(), &|| Live::Frag).unwrap()
  }

  fn virt(s: &str) -> V {
    parse(&mut s.chars(),
          &|spine, arms| Virt::branch(spine, (), arms),
          &|| Virt::Frag(tree("alt()"))).unwrap()
  }

  #[test]
  fn test_parse() {
    assert_eq!(tree("3(1()2())"),
               StrT { spine: "3".into(),
                      children: vec![
                        StrT { spine: "1".into(), children: vec![] },
                        StrT { spine: "2".into(), children: vec![] },
                      ] });
    assert_eq!(live("3(1()*)"),
               Branch { spine: "3".into(),
                        arms: vec![
                          Branch { spine: "1".into(), arms: vec![] }.into(),
                          Live::Frag,
                        ] }.into());
  }

  impl Index<usize> for StrT {
    type Output = StrT;
    fn index(&self, i: usize) -> &StrT {
      &self.children[i]
    }
  }

  impl IndexMut<usize> for StrT {
    fn index_mut(&mut self, i: usize) -> &mut StrT {
      &mut self.children[i]
    }
  }

  impl Default for StrT {
    fn default() -> Self {
      StrT { spine: "".into(), children: vec![] }
    }
  }

  impl Underlying for StrT {
    type Spine = String;
    type Info = ();

    fn create(spine: String, _: (), children: Vec<StrT>) -> Self {
      StrT { spine, children }
    }

    fn set_spine(&mut self, spine: String) {
      self.spine = spine;
    }

    fn set_info(&mut self, _: ()) {}
  }

  impl diff::ApplyPatch<StrT> for StrT {
    fn len(&self) -> usize {
      self.children.len()
    }

    fn remove(&mut self, i: usize) {
      self.children.remove(i);
    }

    fn insert(&mut self, i: usize, c: StrT) {
      self.children.insert(i, c);
    }
  }

  //////////////////////////////////////////////////////////////////////////////////////////////////

  #[test]
  fn test_impose_frag_mut_underlying() {
    let mut u = tree("1()");
    Live::Frag.impose(virt("2()"), &mut u);
    assert_eq!(u, tree("2()"));
  }

  #[test]
  fn test_impose_new_branch() {
    let mut l = live("*");
    let mut u = tree("alt()");
    l.impose(virt("3(1()2())"), &mut u);
    assert_eq!(u, tree("3(1()2())"));
  }

  #[test]
  fn test_impose_update_spine() {
    let mut l = live("1()");
    let mut u = tree("1()");
    l.impose(virt("2()"), &mut u);
    assert_eq!(l, live("2()"));
    assert_eq!(u, tree("2()"));
  }

  #[test]
  fn test_impose_add_arm() {
    let mut l = live("4()");
    let mut u = tree("4()");
    l.impose(Virt::Branch(
      Branch { spine: "4".into(), arms: vec![Virt::Frag(tree("alt()"))] },
      ()
    ), &mut u);
    assert_eq!(l, live("4(*)"));
    assert_eq!(u, tree("4(alt())"));
  }

  #[test]
  fn test_impose_remove_arm() {
    let mut l = live("4(*)");
    let mut u = tree("4(alt())");
    l.impose(virt("4()"), &mut u);
    assert_eq!(l, live("4()"));
    assert_eq!(u, tree("4()"));
  }

  #[test]
  fn test_impose_update_child() {
    let mut l = live("4(3())");
    let mut u = tree("4(3())");
    l.impose(virt("4(2())"), &mut u);
    assert_eq!(l, live("4(2())"));
    assert_eq!(u, tree("4(2())"));
  }

  #[test]
  fn test_impose_rec_update_no_changes() {
    let mut l = live("4(3())");
    let mut u = tree("4(3())");
    l.impose(virt("4(3())"), &mut u);
    assert_eq!(l, live("4(3())"));
    assert_eq!(u, tree("4(3())"));
  }
}
